#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""  
@Project : pythonProject01
@File : ehtools.py
@Author : erhu
@Time : 2025/9/9 18:47  
@脚本说明 : 
"""

from Val.mysql_val import mysql_run
from Val.portscan_val import scan_port_run
from Val.redis_val import redis_run
from Val.ip_val import run_ip
from Val.ccolor import Colors
from Val.ssh_val import ssh_run
from Val.subdomain_val import subdm_run
import argparse
import sys
import os

banner = """
    $$$$$$$$\ $$\   $$\ $$$$$$$$\  $$$$$$\   $$$$$$\  $$\       $$$$$$\  
    $$  _____|$$ |  $$ |\__$$  __|$$  __$$\ $$  __$$\ $$ |     $$  __$$\ 
    $$ |      $$ |  $$ |   $$ |   $$ /  $$ |$$ /  $$ |$$ |     $$ /  \__|
    $$$$$\    $$$$$$$$ |   $$ |   $$ |  $$ |$$ |  $$ |$$ |     \$$$$$$\  
    $$  __|   $$  __$$ |   $$ |   $$ |  $$ |$$ |  $$ |$$ |      \____$$\ 
    $$ |      $$ |  $$ |   $$ |   $$ |  $$ |$$ |  $$ |$$ |     $$\   $$ |
    $$$$$$$$\ $$ |  $$ |   $$ |    $$$$$$  | $$$$$$  |$$$$$$$$\\$$$$$$  |
    \________|\__|  \__|   \__|    \______/  \______/ \________|\______/ 
"""

def print_menu():
    menu = f"""
{Colors.CYAN_BRIGHT}{Colors.BOLD}╔══════════════════════════════════════════════════════════════════════════════╗{Colors.RESET}
{Colors.CYAN_BRIGHT}{Colors.BOLD}║{Colors.RESET}                                  功能菜单                                    {Colors.CYAN_BRIGHT}{Colors.BOLD}║{Colors.RESET}
{Colors.CYAN_BRIGHT}{Colors.BOLD}╠══════════════════════════════════════════════════════════════════════════════╣{Colors.RESET}
{Colors.CYAN_BRIGHT}{Colors.BOLD}║{Colors.RESET}  1. MySQL相关漏扫         4. 子域名扫描          7. 退出程序                 {Colors.CYAN_BRIGHT}{Colors.BOLD}║{Colors.RESET}
{Colors.CYAN_BRIGHT}{Colors.BOLD}║{Colors.RESET}  2. Redis相关漏扫         5. 主机存活探测                                    {Colors.CYAN_BRIGHT}{Colors.BOLD}║{Colors.RESET}
{Colors.CYAN_BRIGHT}{Colors.BOLD}║{Colors.RESET}  3. SSH相关漏扫           6. 全端口扫描                                      {Colors.CYAN_BRIGHT}{Colors.BOLD}║{Colors.RESET}
{Colors.CYAN_BRIGHT}{Colors.BOLD}╚══════════════════════════════════════════════════════════════════════════════╝{Colors.RESET}
    """
    print(f"{Colors.RED_BRIGHT}{Colors.BOLD}{banner}{Colors.RESET}")
    print(menu)


def get_file_path(prompt, default_path):

    while True:
        file_path = input(f"{Colors.CYAN_BRIGHT}{prompt} (默认: {default_path}): {Colors.RESET}") or default_path
        if os.path.exists(file_path):
            return file_path
        else:
            print(f"{Colors.RED_BRIGHT}文件不存在，请重新输入或按回车使用默认路径。{Colors.RESET}")


def interactive_mode():
    while True:
        print_menu()
        try:
            choice = input(f"{Colors.YELLOW_BRIGHT}请选择功能 (1-7): {Colors.RESET}").strip()

            if choice == '1':
                print(f"{Colors.GREEN_BRIGHT}MySQL相关漏扫{Colors.RESET}")
                host = input(
                    f"{Colors.CYAN_BRIGHT}请输入MySQL主机地址 (默认: 127.0.0.1): {Colors.RESET}") or "127.0.0.1"
                port = int(input(f"{Colors.CYAN_BRIGHT}请输入MySQL端口号 (默认: 3306): {Colors.RESET}") or "3306")
                # 自定义字典文件路径
                user_file_path = get_file_path("请输入用户名字典文件路径", "./dir/users.txt")
                pass_file_path = get_file_path("请输入密码字典文件路径", "./dir/passwd.txt")

                with open(user_file_path, "r", encoding="utf-8") as user_file:
                    usernames = user_file.readlines()
                with open(pass_file_path, "r", encoding="utf-8") as pass_file:
                    passwords = pass_file.readlines()
                mysql_run(host, port, usernames, passwords)

            elif choice == '2':
                print(f"{Colors.GREEN_BRIGHT}Redis相关漏扫{Colors.RESET}")
                host = input(
                    f"{Colors.CYAN_BRIGHT}请输入Redis主机地址 (默认: 127.0.0.1): {Colors.RESET}") or "127.0.0.1"
                port = int(input(f"{Colors.CYAN_BRIGHT}请输入Redis端口号 (默认: 6379): {Colors.RESET}") or "6379")
                pass_file_path = get_file_path("请输入密码字典文件路径", "./dir/passwd.txt")

                with open(pass_file_path, "r", encoding="utf-8") as pass_file:
                    passwords = pass_file.readlines()
                redis_run(host, port, passwords)

            elif choice == '3':
                print(f"{Colors.GREEN_BRIGHT}SSH相关漏扫{Colors.RESET}")
                host = input(f"{Colors.CYAN_BRIGHT}请输入SSH主机地址 (默认: 127.0.0.1): {Colors.RESET}") or "127.0.0.1"
                port = int(input(f"{Colors.CYAN_BRIGHT}请输入SSH端口号 (默认: 22): {Colors.RESET}") or "22")
                user_file_path = get_file_path("请输入用户名字典文件路径", "./dir/users.txt")
                pass_file_path = get_file_path("请输入密码字典文件路径", "./dir/passwd.txt")

                with open(user_file_path, "r", encoding="utf-8") as user_file:
                    usernames = user_file.readlines()
                with open(pass_file_path, "r", encoding="utf-8") as pass_file:
                    passwords = pass_file.readlines()
                ssh_run(host, port, usernames, passwords)

            elif choice == '4':
                print(f"{Colors.GREEN_BRIGHT}子域名扫描{Colors.RESET}")
                domain = input(
                    f"{Colors.CYAN_BRIGHT}请输入主域名地址 (默认: example.com): {Colors.RESET}") or "example.com"
                sub_file_path = get_file_path("请输入子域名字典文件路径", "./dir/sub.txt")

                with open(sub_file_path, "r", encoding="utf-8") as sub_file:
                    subs = sub_file.readlines()
                subdm_run(domain, subs)

            elif choice == '5':
                print(f"{Colors.GREEN_BRIGHT}主机存活探测{Colors.RESET}")
                host = input(
                    f"{Colors.CYAN_BRIGHT}请输入IP主机地址网段 (默认: 127.0.0.1): {Colors.RESET}") or "127.0.0.1"
                iface = input(f"{Colors.CYAN_BRIGHT}请输入扫描网卡 (默认: 以太网): {Colors.RESET}") or "以太网"
                run_ip(host, iface)

            elif choice == '6':
                print(f"{Colors.GREEN_BRIGHT}全端口扫描{Colors.RESET}")
                host = input(f"{Colors.CYAN_BRIGHT}请输入目标主机地址 (默认: 127.0.0.1): {Colors.RESET}") or "127.0.0.1"
                sep = int(input(f"{Colors.CYAN_BRIGHT}请输入单线程扫描端口个数 (默认: 10): {Colors.RESET}") or "10")
                scan_port_run(host, sep)

            elif choice == '7':
                print(f"{Colors.GREEN_BRIGHT}感谢使用EHTOOLS，再见！{Colors.RESET}")
                sys.exit(0)

            else:
                print(f"{Colors.RED_BRIGHT}无效选择，请输入1-7之间的数字。{Colors.RESET}")

            input(f"\n{Colors.YELLOW_BRIGHT}按回车键继续...{Colors.RESET}")

        except KeyboardInterrupt:
            print(f"\n\n{Colors.GREEN_BRIGHT}感谢使用EHTOOLS，再见！{Colors.RESET}")
            sys.exit(0)
        except Exception as e:
            print(f"{Colors.RED_BRIGHT}发生错误: {str(e)}{Colors.RESET}")
            input(f"{Colors.YELLOW_BRIGHT}按回车键继续...{Colors.RESET}")


# 创建命令行解析器对象
parser = argparse.ArgumentParser(description="EHTOOLS内网工具")
subparsers = parser.add_subparsers(help="请选择要使用的功能\n", dest="sub_command")
# 创建MySQL的解析器对象
mysql_subparsers = subparsers.add_parser("mysql", help="MySQL相关漏扫")
mysql_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="MySQL主机地址")
mysql_subparsers.add_argument("--port", "-P", dest="port", type=int, default=3306, help="MySQL端口号")
mysql_subparsers.add_argument("--username", "-u", dest="username", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dir/users.txt", help="MySQL用户名")
mysql_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dir/passwd.txt", help="MySQL密码本")
# 创建Redis的解析器对象
redis_subparsers = subparsers.add_parser("redis", help="Redis相关漏扫")
redis_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="Redis主机地址")
redis_subparsers.add_argument("--port", "-P", dest="port", type=int, default=6379, help="Redis端口号")
redis_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dir/passwd.txt", help="Redis密码本")

ssh_subparsers = subparsers.add_parser("ssh", help="SSH相关漏扫")
ssh_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="SSH主机地址")
ssh_subparsers.add_argument("--port", "-P", dest="port", type=int, default=22, help="SSH端口号")
ssh_subparsers.add_argument("--username", "-u", dest="username", type=argparse.FileType(mode='r', encoding="utf-8"),
                            default="./dir/users.txt", help="SSH用户名")
ssh_subparsers.add_argument("--password", "-p", dest="password", type=argparse.FileType(mode='r', encoding="utf-8"),
                            default="./dir/passwd.txt", help="SSH密码本")

subdm_subparsers = subparsers.add_parser("subdm", help="子域名相关漏扫")
subdm_subparsers.add_argument("--domain", "-D", dest="domain", type=str, default="example.com", help="主域名地址")
subdm_subparsers.add_argument("--sub", "-s", dest="sub", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dir/sub.txt", help="子域名合集")

# 创建IP的解析器对象
ip_subparsers = subparsers.add_parser("ip", help="主机相关漏扫")
ip_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="IP主机地址网段")
ip_subparsers.add_argument("--iface", "-f", dest="iface", type=str, default="以太网", help="扫描网卡")

portscan_subparsers = subparsers.add_parser("portscan", help="端口扫描")
portscan_subparsers.add_argument("--host", "-H", dest="host", type=str, default="127.0.0.1", help="目标主机地址")
portscan_subparsers.add_argument("--sep", "-S", dest="sep", type=int, default=10, help="单线程扫描端口个数")

args = parser.parse_args()
if args.sub_command == "mysql":
    mysql_run(args.host, args.port, args.username.readlines(), args.password.readlines())
elif args.sub_command == "redis":
    redis_run(args.host, args.port, args.password.readlines())
elif args.sub_command == "ssh":
    ssh_run(args.host, args.port, args.username.readlines(), args.password.readlines())
elif args.sub_command == "subdm":
    subdm_run(args.domain, args.sub.readlines())
elif args.sub_command == "ip":
    run_ip(args.host, args.iface)
elif args.sub_command == "portscan":
    scan_port_run(args.host, args.sep)
else:
    # 如果没有提供命令行参数，则进入交互模式
    interactive_mode()
